package com.ruoyi.ai.splitter;

import org.springframework.ai.document.Document;
import org.springframework.ai.transformer.splitter.TextSplitter;

import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class MarkdownSplitter extends TextSplitter {

    // 正则表达式用于匹配二级标题（## 开头的行）
    private static final Pattern HEADING_PATTERN = Pattern.compile("^## (.+)$", Pattern.MULTILINE);
    

    @Override
    protected List<String> splitText(String text) {
       return Collections.emptyList();
    }
    
    /**
     * 重写split方法，实现按二级标题分割文档的逻辑
     */
    @Override
    public List<Document> split(Document document) {
        String text = document.getText();
        List<Document> documents = new ArrayList<>();
        
        // 查找所有二级标题
        Matcher matcher = HEADING_PATTERN.matcher(text);
        
        // 存储所有标题的位置和内容
        List<Integer> headingPositions = new ArrayList<>();
        List<String> headingTexts = new ArrayList<>();
        
        while (matcher.find()) {
            headingPositions.add(matcher.start());
            headingTexts.add(matcher.group(1)); // 提取标题文本（不包括##）
        }
        
        // 如果没有找到二级标题，返回原始文档
        if (headingPositions.isEmpty()) {
            return List.of(document);
        }
        
        // 处理每个标题及其内容
        for (int i = 0; i < headingPositions.size(); i++) {
            int startPos = headingPositions.get(i);
            int endPos;
            
            // 如果是最后一个标题，则内容一直到文档结束
            if (i == headingPositions.size() - 1) {
                endPos = text.length();
            } else {
                // 否则，内容一直到下一个标题开始
                endPos = headingPositions.get(i + 1);
            }
            
            // 提取标题下的内容（不包括标题本身）
            String titleText = headingTexts.get(i);
            //将titleText转成 什么是titleText，titleText的理解，titleText如何使用，titleText的优点，titleText有什么用，说说你对titleText的理解，体系介绍titleText
            titleText = "什么是" + titleText + "，" + titleText + "的理解，" + titleText + "的优点，" + titleText + "有什么用，" + titleText + "的体系介绍";
            String contentText = text.substring(startPos, endPos).trim();
            
            // 从内容中移除标题行
            String contentWithoutTitle = contentText.replaceFirst("^## " + Pattern.quote(titleText) + "\\R?", "").trim();
            
            // 创建metadata，包含标题下的内容
            Map<String, Object> metadata = new HashMap<>(document.getMetadata()); // 保留原始metadata
            metadata.put("content", contentWithoutTitle);
            
            // 创建新的Document，以标题为text，内容为metadata
            Document newDoc = new Document(titleText, metadata);
            documents.add(newDoc);
        }
        
        return documents;
    }
    
    /**
     * 测试主程序
     */
    public static void main(String[] args) {
        // 示例Markdown文本
        String markdownText = "# 文档标题\n\n" +
                             "这是一级标题下的内容。\n\n" +
                             "## 第一个二级标题\n\n" +
                             "这是第一个二级标题下的内容。\n" +
                             "包含多行文本。\n\n" +
                             "### 三级标题\n\n" +
                             "这是三级标题下的内容。\n\n" +
                             "## 第二个二级标题\n\n" +
                             "这是第二个二级标题下的内容。\n\n" +
                             "## 第三个二级标题\n\n" +
                             "这是最后一个二级标题下的内容。";

        // 创建 MarkdownSplitter 实例
        MarkdownSplitter splitter = new MarkdownSplitter();
        
        // 模拟 Document 输入
        Document document = new Document(markdownText, Map.of("source", "example.md"));
        
        // 调用 split 方法
        List<Document> result = splitter.split(document);
        
        // 打印分割结果
        for (int i = 0; i < result.size(); i++) {
            Document doc = result.get(i);
            System.out.println("Document " + (i + 1) + ":");
            System.out.println("Title: " + doc.getText());
            System.out.println("Content: " + doc.getMetadata().get("content"));
            System.out.println();
        }
    }
}